# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/5 10:38 
@Author  : ZhangShenao 
@File    : base.py
@Desc    : 利用子图架构实现多Agent协助的营销智能体

子图是一个非常有用的设计原则
对于一个较为复杂的大型应用,可以将整个应用拆分成多个模块
每个模块独立封装成一个子图,每个子图各自管理自己的执行流程和状态数据
最后将多个子图组合成一个大型复杂应用

因为LangGraph的节点可以是任意的Python函数或Runnable可运行组件
而Graph也可以编译成Runnable可运行组件
因此,可以将父图中的一个节点设置为子图的Graph
这样利用LangGraph的子图,可以轻松实现类似多Agent亦或者AI工作流之类的功能

创建子图架构,最核心的部分就是要管理节点之间的状态传递
所有子图会共享父图的状态,并且每个子图也可以维护自己的私有状态
子图的私有状态之间是互相隔离的

需要特别注意的是,需要为共享状态中的每个字段都定义归纳函数
因为多个子图之间是并行执行的,无法保证执行的先后顺序
如果不为共享状态定义归纳函数,则后执行完成的子图可能会把先执行完成的子图的状态覆盖掉
实际上在LangGraph编译的过程中,如果一个图有多边并行的情况,并且没有为每个字段都定义归纳函数,则会直接抛出异常

使用多图架构的优势:
1. 模块化管理: 每个子图(Agent)可以独立维护各自的状态,而不需要在公共状态中考虑所有Agents
2. 可扩展性: 后续如果想要增加新的Agent,只需要在父图添加一个节点即可,无需改动原有结构
"""
import os
from typing import TypedDict, Annotated

import dotenv
from langchain_core.language_models import BaseChatModel
from langchain_openai import ChatOpenAI


# 定义状态字段的归纳函数
def reduce_str(left: str | None, right: str | None) -> str:
    """
    状态字段的归纳函数
    :param left: 当前的状态字符串
    :param right: 更新的状态字符串
    :return: 更新后的状态字符串
    """

    # 只有当更新的状态字符串不为空时,才更新状态,避免状态间的互相覆盖
    if right is not None and right != "":
        return right
    return left


# 定义父图的公共状态
class BaseAgentState(TypedDict):
    """多智能体的公共状态"""

    query: Annotated[str, reduce_str]
    """用户的原始提问"""

    live_content: Annotated[str, reduce_str]
    """直播带货营销文案"""

    xhs_content: Annotated[str, reduce_str]
    """小红书营销文案"""


def get_llm() -> BaseChatModel:
    """获取LLM模型"""
    
    dotenv.load_dotenv()
    return ChatOpenAI(model="gpt-4o-mini", temperature=0, openai_api_base=os.getenv("OPENAI_API_BASE"))
